import numpy as np
import random
from PIL import Image
import torchvision.transforms as transforms
import imgaug.augmenters as iaa
import torch


class AugMix:
    def __init__(self, preprocess, severity=3, width=3, depth=-1, alpha=1.):
        self.preprocess = preprocess
        self.severity = severity
        self.width = width
        self.depth = depth
        self.alpha = alpha
        self.augmentations = [
            iaa.GaussianBlur((0.1, 3.0)),
            iaa.AdditiveGaussianNoise(scale=(0, 0.2 * 255)),
            iaa.Posterize(nb_bits=(4, 8)),
            iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 2.0)),
            iaa.Multiply((0.4, 1.6)),
            iaa.ContrastNormalization((0.5, 1.5)),
            iaa.Grayscale(alpha=(0.0, 1.0)),
            iaa.Affine(rotate=(-25, 25)),
        ]

    def __call__(self, image):
        ws = np.float32(np.random.dirichlet([self.alpha] * self.width))
        m = np.float32(np.random.beta(self.alpha, self.alpha))

        mix = torch.zeros_like(self.preprocess(image))
        for i in range(self.width):
            image_aug = image.copy()
            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
            for _ in range(depth):
                op = random.choice(self.augmentations)
                image_aug = Image.fromarray(op(image=np.asarray(image_aug)))
            mix += ws[i] * self.preprocess(image_aug)

        mixed = (1 - m) * self.preprocess(image) + m * mix
        return mixed



def img_param_init(args):
    dataset = args.dataset
    if dataset == 'Office31':
        domains = ['amazon', 'dslr', 'webcam']
    elif dataset == 'OfficeHome':
        domains = ['Art', 'Clipart', 'Product', 'Real_World']
    elif dataset == 'Visda17':
        domains = ['train', 'validation']
    elif dataset == 'DomainNet':
        domains = ["clipart", "infograph", "painting", "quickdraw", "real", "sketch"]
    else:
        print('No such dataset exists!')
    
    args.domains = domains
    args.img_dataset = {
        'Office31': ['amazon', 'dslr', 'webcam'],
        'OfficeHome': ['Art', 'Clipart', 'Product', 'Real_World'],
        'Visda17': ['train', 'validation'],
        'DomainNet': ["clipart", "infograph", "painting", "quickdraw", "real", "sketch"]
    }
    args.input_shape = (3, 224, 224)
    
    if args.dataset == 'Office31':
        args.num_classes = 31
    elif args.dataset == 'OfficeHome':
        args.num_classes = 65
    elif args.dataset == 'Visda17':
        args.num_classes = 12
    elif args.dataset == 'DomainNet':
        args.num_classes = 345
        
    return args